/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv.region;

import com.google.protobuf.ByteString;
import com.google.protobuf.Descriptors;
import com.webank.tikv.TiConfiguration;
import com.webank.tikv.exception.GrpcException;
import com.webank.tikv.exception.InternalException;
import com.webank.tikv.exception.RetryableException;
import com.webank.tikv.exception.TiKVException;
import com.webank.tikv.kvproto.Errorpb;
import com.webank.tikv.kvproto.Kvrpcpb;
import com.webank.tikv.kvproto.Metapb;
import com.webank.tikv.kvproto.TikvGrpc;
import com.webank.tikv.operation.RegionErrorHandler;
import com.webank.tikv.util.BackOffer;
import com.webank.tikv.util.ChannelFactory;
import io.grpc.ManagedChannel;
import io.grpc.stub.StreamObserver;
import org.apache.commons.lang3.tuple.Triple;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
import java.util.function.Supplier;

/**
 * Note that RegionStoreClient itself is not thread-safe
 */
public class RegionStoreClient extends AbstractRegionStoreClient {

    private final RegionErrorHandler<Kvrpcpb.RedisResponse> regionErrorHandler;

    private RegionStoreClient(TiConfiguration conf, TiRegion region, TiStore store,
                              ChannelFactory channelFactory,
                              TikvGrpc.TikvBlockingStub blockingStub, TikvGrpc.TikvFutureStub asyncStub,
                              RegionManager regionManager) {
        super(conf, region, store, channelFactory, blockingStub, asyncStub, regionManager);
        this.regionErrorHandler = new RegionErrorHandler<>(regionManager, this,
            resp -> resp.hasRegionError() ? resp.getRegionError() : null);
    }

    public Kvrpcpb.RedisResponse execRedis(BackOffer backOffer, String cmd, ByteString key, ByteString... value) {
        Supplier<Kvrpcpb.RedisRequest> factory = buildRedisRequestFactory(
            region.makeLeaderContext(), System.currentTimeMillis(), cmd, key, value);
        Kvrpcpb.RedisResponse resp = null;
        for (int i = 0; i < 3; i++) {
            resp = callWithRetry(backOffer, TikvGrpc.getRedisExecMethod(), factory, regionErrorHandler);
            if (resp != null && !resp.hasRegionError()) {
                break;
            } else if (regionErrorHandler.handleResponseError(backOffer, resp)) {
                continue;
            }
            break;
        }
        if (resp == null) {
            throw new InternalException("response failed without any cause");
        }
        return postRedisResponse(resp);
    }

    public CompletableFuture<Kvrpcpb.RedisResponse> execRedisAsync(BackOffer backOffer,
                                                                   ScheduledExecutorService executor,
                                                                   String cmd, ByteString key, ByteString... value) {
        return execRedisAsync(region.makeLeaderContext(), backOffer, executor,
            System.currentTimeMillis(), cmd, key, value);
    }

    public CompletableFuture<Kvrpcpb.RedisResponse> execRedisAsync(Kvrpcpb.Context context, BackOffer backOffer,
                                                                   ScheduledExecutorService executor,
                                                                   long sendTime, String cmd,
                                                                   ByteString key, ByteString... value) {
        CompletableFuture<Kvrpcpb.RedisResponse> future = new CompletableFuture<>();
        Supplier<Kvrpcpb.RedisRequest> factory = buildRedisRequestFactory(context, sendTime, cmd, key, value);
        fireRedisRequest(backOffer, executor, factory, future);
        return future;
    }

    private void fireRedisRequest(BackOffer backOffer, ScheduledExecutorService executor,
                                  Supplier<Kvrpcpb.RedisRequest> factory,
                                  CompletableFuture<Kvrpcpb.RedisResponse> future) {
        callAsyncWithRetry(backOffer, executor, TikvGrpc.getRedisExecMethod(), factory,
            new StreamObserver<Kvrpcpb.RedisResponse>() {
                @Override
                public void onNext(Kvrpcpb.RedisResponse response) {
                    if (response.hasRegionError()) {
                        long time = regionErrorHandler.calNextRetryTime(backOffer, response);
                        if (time >= 0) {
                            future.completeExceptionally(new RetryableException(time));
                        } else {
                            String msg = extractRegionErrorMsg(response.getRegionError());
                            future.completeExceptionally(new TiKVException(msg));
                        }
                        return;
                    }
                    future.complete(postRedisResponse(response));
                }

                @Override
                public void onError(Throwable t) {
                    future.completeExceptionally(t);
                }

                @Override
                public void onCompleted() {
                }
            }, regionErrorHandler);
    }

    // 构建redis request
    private Supplier<Kvrpcpb.RedisRequest> buildRedisRequestFactory(Kvrpcpb.Context context, long sendTime,
                                                                    String cmd, ByteString key, ByteString... value) {
        return () -> {
            Kvrpcpb.RedisRequest.Builder builder = Kvrpcpb.RedisRequest.newBuilder()
                .setContext(context)
                .setCmd(cmd)
                .setKey(key);
            for (ByteString bytes : value) {
                builder.addArgs(bytes);
            }

            Kvrpcpb.RTTDetail reqRttDetail = Kvrpcpb.RTTDetail.newBuilder()
                .setSendTimeMs(sendTime)
                .build();
            builder.setRttDetail(reqRttDetail);
            return builder.build();
        };
    }

    // 记录收到回包时间
    private Kvrpcpb.RedisResponse postRedisResponse(Kvrpcpb.RedisResponse response) {
        Kvrpcpb.RTTDetail respRttDetail = response.getRttDetail().toBuilder()
            .setReceiveTimeMs(System.currentTimeMillis())
            .build();
        return response.toBuilder().setRttDetail(respRttDetail).build();
    }

    private String extractRegionErrorMsg(Errorpb.Error regionError) {
        String msg = regionError.getMessage();
        if (!msg.isEmpty()) {
            return msg;
        }
        for (Map.Entry<Descriptors.FieldDescriptor, Object> entry : regionError.getAllFields().entrySet()) {
            if (entry.getKey().hasPresence()) {
                return String.format("%s, %s", entry.getKey(), entry.getValue());
            }
        }
        return "can not extract err msg";
    }

    public static class RegionStoreClientBuilder {

        private final TiConfiguration conf;
        private final ChannelFactory channelFactory;
        private final RegionManager regionManager;

        public RegionStoreClientBuilder(TiConfiguration conf, ChannelFactory channelFactory,
                                        RegionManager regionManager) {
            Objects.requireNonNull(conf, "conf is null");
            Objects.requireNonNull(channelFactory, "channelFactory is null");
            Objects.requireNonNull(regionManager, "regionManager is null");
            this.conf = conf;
            this.channelFactory = channelFactory;
            this.regionManager = regionManager;
        }

        public RegionStoreClient build(TiRegion region, TiStore store) throws GrpcException {
            Objects.requireNonNull(region, "region is null");
            Objects.requireNonNull(store, "store is null");

            ManagedChannel channel = channelFactory.getChannel(store.getStore().getAddress());
            TikvGrpc.TikvBlockingStub blockingStub = TikvGrpc.newBlockingStub(channel);
            TikvGrpc.TikvFutureStub asyncStub = TikvGrpc.newFutureStub(channel);

            // todo cache by store and region
            return new RegionStoreClient(conf, region, store, channelFactory,
                blockingStub, asyncStub, regionManager);
        }

        public RegionStoreClient build(ByteString key, BackOffer backOffer) throws GrpcException {
            Triple<TiRegion, TiStore, Metapb.Peer> triple =
                regionManager.getRegionStoreByKey(key, backOffer, false);
            return build(triple.getLeft(), triple.getMiddle());
        }
    }

}
